'''
Author: your name
Date: 2022-12-21 16:14:11
LastEditTime: 2022-12-23 13:09:11
LastEditors: LAPTOP-4PHAO8B6
Description: In User Settings Edit
FilePath: \dy-pro\model.py
'''

import whisper
from whisper.utils import *
import os

from splitvideo import *

os.environ["CUDA_VISIBLE_DEVICES"] = "0"

def model(data):

    file_path  = data["file_path"]
    save_path  = data["save_path"]
    language   = data["language"]
    cpu_or_gpu = data["s_cpu"]
    model_name = data["model"]
    file_type  = data["file_type"]
    file_name  = data["file_name"]

    s = file_name.split('.')[0]

    if file_type != "mp3":
        file_path = split_video(file_path,file_name,"mp3")

    try:
        if cpu_or_gpu == "gpu" and model_name != "medium":

            model = whisper.load_model(model_name,"cuda")
        else:
            model = whisper.load_model(model_name)



        result = model.transcribe(file_path,language=language)



        with open(save_path+'/'+s+".srt", "w+", encoding="utf-8") as f:
            write_srt(result["segments"],file=f)
    except Exception as e:
        print(e)

    return True

